Could you help me understand why this approach was chosen ?
I can make some sense of 512 thread groups being the target (if the M or N are larger, we go to the regular matmul, if the K is larger still we just loop more inside the kernel), but would like to know how you ended up at that number

Also, we do employ checks such that K is always divisible by group_size - so wouldn't the looping step just then come down to be equivalent to split_k = std::max(split_k, K / group_size)

For the 512 target, it was mostly an empirical choice trying to balance GPU occupancy with atomic reduction overhead. My mental math was that on the larger Apple Silicon chips (like the 76-core M2 Ultra), we generally need a handful of active threadgroups per core to effectively hide memory latency (roughly $76 \times 6 \approx 456$). So, aiming for ~512 felt like a safe baseline to ensure good saturation even on the high-end chips. On the flip side, I wanted to avoid pushing the split factor too high to keep the atomic add contention in check during the final accumulation, as over-splitting the K-dimension can quickly offset the occupancy gains.Ideally, this target should be a dynamic function of the active GPU core count. However, since dynamically querying the exact core count in Metal isn't straightforward, tuning for the high-end ceiling acts as a safe catch-all. The hardware scheduler on smaller chips (like the base M1/M2) handles the oversubscription of 512 threadgroups gracefully without much penalty, whereas under-subscribing an Ultra would be catastrophic for latency hiding. If MLX has a preferred internal way to dynamically scale split-K factors based on device tiers, I'm completely open to refactoring this!

I originally avoided just using split_k = std::min(split_k, K / group_size) because we need to guarantee quantization group alignment, not just an upper bound.

If we strictly use std::min, it might result in a split_k that doesn't perfectly divide K into multiples of group_size. For example, if K=1024, group_size=64, and the calculated split_k is 12, the chunk size 1024 / 12 does not align with 64(the quant group). If a threadgroup's chunk starts unaligned with group_size, it will read misaligned scale and bias values, leading to incorrect numerical results.

The while loop ensures we step down until (K / split_k) % group_size == 0 so that every threadgroup boundary perfectly aligns with the quantization blocks. Let me know if you think there is a more elegant way to enforce this modulo alignment mathematically instead of the loop!

Good point

I think the way around this would be support unaligned K_eff for the last partition (that logic doesn't exist in the QMMs but does in the regular MM code) - we certainly don't want it to be the case that you end up with a K that divides to 31 quantized groups, but doesn't end up dispatched to the split-K variant because 31 is a prime and we require the number of splits to perfectly divide K / group_size

That said, we can merge this for now and do that in a follow up PR

For now, just so we are safe, could you add a check here to make sure that K is large enough (in relation to M and N) to warrant going through the loop of splits ?
We naturally short-circuit if n_tiles * m_tiles >= 512 so large M and N are covered - it would be good to similarly short circuit out if the K isn't too large compared to the M and N.

Till we have a fix for the loop, I would like to avoid to a 64x64x128 matmul as an example, going through a 100 iterations where it could have just short circuited earlier

After that, we should be ready to merge!